#include "CommandLineFilteredTraceEventCallback.h"

HRESULT STDMETHODCALLTYPE CommandLineFilteredTraceEventCallback::OnEvent(_In_ ITraceEvent* traceEvent, _In_ ITraceRelogger* relogger)
{
    PEVENT_RECORD eventRecord;
    SUCCEEDED(traceEvent->GetEventRecord(&eventRecord));

    PEVENT_HEADER eventHeader = &eventRecord->EventHeader;
    GUID providerGuid = eventHeader->ProviderId;
    int pid = eventHeader->ProcessId;

    /* Process of Interest Events */
    if (pid == this->commandLineToProcessId)
    {
        return relogger->Inject(traceEvent);
    }

    /* System-wide Process and Image Load Events */
    if (providerGuid == ProcessEventGuid)
    {
        switch (eventHeader->EventDescriptor.Opcode)
        {
        case 3:
            this->ParseAndSetProcessInformation(eventRecord);
            return relogger->Inject(traceEvent);
        case 1:
            if (this->ParseAndSetProcessInformation(eventRecord))
            {
                return relogger->Inject(traceEvent);
            }
            return S_OK;
        }
    }

    /* Only relog image load/unload events for process of interest */
    if (providerGuid == ImageEventGuid)
    {
        switch (eventHeader->EventDescriptor.Opcode)
        {
        case 2:
        case 10:
            CopyMemory(&pid, (PBYTE)eventRecord->UserData + 16, sizeof(int));
            if (pid == this->commandLineToProcessId)
            {
                return relogger->Inject(traceEvent);
            }

            return S_OK;
        case 3:
            return relogger->Inject(traceEvent);
        }
    }

    return S_OK;
}

bool CommandLineFilteredTraceEventCallback::ParseAndSetProcessInformation(PEVENT_RECORD eventRecord)
{
    PBYTE userData = (PBYTE)eventRecord->UserData;
    int pid = -1;
    int userDataLength = eventRecord->UserDataLength;

    int sid = 0, imageNameIndex = 0;
    CopyMemory(&sid, userData + VERSION4OFFSET, sizeof(int));
    if (sid == 0)
    {
        imageNameIndex = VERSION4OFFSET + 4;
    }
    else
    {
        BYTE numAuthorities;
        CopyMemory(&numAuthorities, userData + VERSION4OFFSET + (TOKENSIZE + 1), sizeof(BYTE));
        imageNameIndex = VERSION4OFFSET + TOKENSIZE + 8 + 4 * numAuthorities;
    }

    BYTE null_ptr;
    while (true)
    {
        CopyMemory(&null_ptr, userData + imageNameIndex, sizeof(BYTE));
        imageNameIndex++;
        if (null_ptr == 0)
        {
            break;
        }
    }

    int charactersCount = (userDataLength + 2) / 2;
    WCHAR* pwsz = (WCHAR*)malloc(userDataLength + 2); // + 2 for wide '\0'

    int commandLineIndex = imageNameIndex;
    WCHAR charbychar;
    int numChars = 0;

    while (true)
    {
        CopyMemory(&charbychar, userData + commandLineIndex, 2);
        commandLineIndex += 2;

        pwsz[numChars++] = charbychar;

        if (charbychar == 0)
        {
            break;
        }
    }

    std::wstring cmdLine(pwsz);

    free(pwsz);

    std::wsmatch match;
    std::regex_search(cmdLine, match, this->commandLineRegex);

    if (!match.empty())
    {
        CopyMemory(&this->commandLineToProcessId, userData + 8, sizeof(int));
        return true;
    }

    return false;
}